/*
 * Alignment.cpp
 *
 *  Created on: Sep 10, 2012
 *      Author: pickrell
 */
#include "Alignment.h"
using namespace std;

Alignment::Alignment(){

}

Alignment::Alignment(string infile){
	read_alignment(infile);
	nchrom = aln.size();
	nsite = aln[0].size();
	nmono = len - nsite;
	compress();
}

void Alignment::read_alignment(string infile){
	aln.clear();
	names.clear();
	igzstream in(infile.c_str());
	vector<string> line;
	struct stat stFileInfo;
	int intStat;
	intStat = stat(infile.c_str(), &stFileInfo);
	if (intStat !=0){
         std::cerr<< "ERROR: cannot open file " << infile << "\n";
         exit(1);
	}
	string st;
	int index = 0;
	while(getline(in, st)){
		string buf;
		stringstream ss(st);
		line.clear();
		while (ss>> buf){
			line.push_back(buf);
		}
		string tmpstr = line[0];
		if (tmpstr.substr(0, 1) == "#") {
			len = atoi(line[1].c_str());
			continue;
		}
		names.push_back(tmpstr);
		name2index.insert(make_pair(tmpstr, index));
		tmpstr = line[1];

		vector<bool> tmp;
		for (int i = 0; i < tmpstr.size(); i++){
			if (tmpstr.substr(i,1) == "0") tmp.push_back(false);
			else if (tmpstr.substr(i, 1) == "1") tmp.push_back(true);
			else {
				cerr << "ERROR: NOT A VALID ALLELE: "<< tmpstr[i] << "\n";
				exit(1);
			}
		}
		index++;
		aln.push_back(tmp);
	}
}

void Alignment::compress(){
	colcounts.clear();
	map<string, int> s2int;
	vector<vector<bool> > naln;
	for (int i = 0; i < nchrom; i++){
		vector<bool> tmp;
		naln.push_back(tmp);
	}
	for (int i = 0; i < nsite; i++){
		string t = "";
		for (int j = 0; j < nchrom; j++){
			bool test = aln[j][i];
			if (test) t += "1";
			else  t += "0";
		}
		if (s2int.find(t) == s2int.end()){
			colcounts.push_back(1);
			s2int.insert(make_pair(t, colcounts.size()-1));
			for (int j = 0; j < nchrom ; j++){
				bool test  = aln[j][i];
				naln[j].push_back(test);
			}
		}
		else colcounts[ s2int[t] ] ++;
	}
	aln = naln;
	nsite = aln[0].size();
}
void Alignment::print(string outfile){
	ogzstream out(outfile.c_str());
	out << "#len "<< len << "\n";
	for (int i = 0; i < nchrom ; i++){
		out << names[i] << " ";
		for (int j = 0; j < nsite ; j ++){
			out << aln[i][j];
		}
		out << "\n";
	}
}

vector<string> Alignment::list_haps(){
	vector<string> toreturn;
	for (vector<string>::iterator it = names.begin(); it != names.end(); it++) toreturn.push_back(*it);
	return toreturn;
}

vector<bool> Alignment::get_hap(string which){
	int w = name2index[which];
	vector<bool> toreturn = aln[w];
	return toreturn;
}

vector<vector<double> > Alignment::get_dists(){
	vector<vector<double> > toreturn;
	for (int i = 0; i < names.size(); i++){
		vector<double> tmp(names.size(), 0);
		toreturn.push_back(tmp);
	}
	for (int i = 0; i < names.size(); i++){
		for (int j = i; j < names.size(); j++){
			double d = get_dist(names[i], names[j]);
			toreturn[i][j] = d;
			toreturn[j][i] = d;
		}
	}
	return toreturn;
}

double Alignment::get_dist(string h1, string h2){
	if (h1 == h2) return 0.0;
	int c1 = name2index[h1];
	int c2 = name2index[h2];
	double toreturn = 0;
	for (int i = 0; i < nsite; i++){
		bool test1 = aln[c1][i];
		bool test2 = aln[c2][i];
		if (test1 != test2){
			int num = colcounts[i];
			toreturn += (double) num;
		}
	}
	toreturn = toreturn/ (double) len;
	return toreturn;
}
